import math
import time
import rclpy
from rclpy.publisher import Publisher
from rclpy.subscription import Subscription
import tf_transformations
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Odometry, Path
from std_msgs.msg import Float64MultiArray
from general_interface.msg import ControlSignal,ChassisServo
from rclpy.qos import QoSProfile, DurabilityPolicy
from message_filters import Subscriber, TimeSynchronizer

from .model import Model
from .mpc_solver import MPC
from .refer_path import Spline, ReferPath


class MpcPathTracker(Node):
    def __init__(self):
        super().__init__("mpc_path_tracker")
        self.get_logger().info("mpc_path_tracker node has been started")
        self.declare_parameter("max_acc", 1.0)
        self.declare_parameter("max_vec", 1.0)
        self.declare_parameter("max_gyro", math.pi / 6)
        self.declare_parameter("mat_q", [1.0, 1.0, 1.0, 1.0])
        self.declare_parameter("mat_r", [1.0, 1.0])
        self.maxAcc: float = self.get_parameter("max_acc").value
        self.maxVec: float = self.get_parameter("max_vec").value
        self.maxGyro: float = self.get_parameter("max_gyro").value
        self.matQ: list[float] = self.get_parameter("mat_q").value
        self.matR: list[float] = self.get_parameter("mat_r").value

        self.ctrlPub: Publisher = self.create_publisher(
            ChassisServo, "control_signal", 10
        )
        self.odomSub: Subscription = self.create_subscription(
            Odometry, "odom", self.odomCallback, 10
        )
        self.pathSub: Subscription = self.create_subscription(
            Path,
            "reference_path",
            self.pathCallback,
            QoSProfile(depth=1, durability=DurabilityPolicy.TRANSIENT_LOCAL),
        )
        self.velocitySub: Subscription = self.create_subscription(
            Float64MultiArray,
            "reference_velocity",
            self.velocityCallback,
            QoSProfile(depth=1, durability=DurabilityPolicy.TRANSIENT_LOCAL),
        )

        # pathSub = Subscriber(self, Path, "reference_path")
        # vecSub = Subscriber(self, Float64MultiArray, "reference_velocity")
        # self.pathSub = TimeSynchronizer([pathSub, vecSub], 10)
        # self.pathSub.registerCallback(self.pathCallback)

        self.predictionPub: Publisher = self.create_publisher(
            Path, "prediction_path", 10
        )
        self.referPathPub: Publisher = self.create_publisher(Path, "mpc_refer", 10)

        self.path = ReferPath(None)
        self.solver = MPC(
            Model(
                [-self.maxVec, self.maxVec],
                [-self.maxAcc, self.maxAcc],
                [-self.maxGyro, self.maxGyro],
            ),
            self.path,
            10,
            0.5,
            4,
            2,
            self.matQ,
            self.matR,
        )
        self.vec = None
        self.positions = None

    def odomCallback(self, msg: Odometry):
        if not self.solver.is_ready():
            # self.get_logger().info("Waiting for path...")
            return

        theta = tf_transformations.euler_from_quaternion(
            [
                msg.pose.pose.orientation.x,
                msg.pose.pose.orientation.y,
                msg.pose.pose.orientation.z,
                msg.pose.pose.orientation.w,
            ]
        )[2]
        state = [
            theta,
            msg.twist.twist.linear.x,
            msg.pose.pose.position.x,
            msg.pose.pose.position.y,
        ]
        refState, refControl = self.solver.path.getReferancePath(
            state, self.solver.T, self.solver.timeStep
        )
        refState = [
            (refState[0, i], refState[1, i], refState[2, i], refState[3, i])
            for i in range(self.solver.T + 1)
        ]
        referPath = self.createPath(refState)
        self.referPathPub.publish(referPath)

        result = self.solver.getControl(state)
        if result is None:
            self.get_logger().info("solve failed")
            return
        ctrl, predStates, predControls = result

        #cs = ControlSignal()
        #cs.forward_control = math.atan(ctrl[1] / 9.8)  # according experiment
        #cs.rotate_control = ctrl[0]
        #cs.timestamp = self.get_clock().now().to_msg()

        cs = ChassisServo()
        cs.angle = math.atan(ctrl[1] / 9.8)
        cs.angle_gain = 1.0
        cs.gyroz = ctrl[0]
        self.ctrlPub.publish(cs)

        predPath = self.createPath(predStates)
        self.predictionPub.publish(predPath)

    def createPath(self, states: list[tuple[float, float, float, float]]) -> Path:
        """
        states is a list of (theta,v, x, y)
        """
        path = Path()
        path.header.stamp = self.get_clock().now().to_msg()
        path.header.frame_id = "map"
        for state in states:
            pose_stamped = PoseStamped()
            pose_stamped.header = path.header
            pose_stamped.pose.position.x = state[2]
            pose_stamped.pose.position.y = state[3]
            pose_stamped.pose.position.z = 0.0  # 假设在二维平面上，z坐标为0

            quaternion = tf_transformations.quaternion_from_euler(0, 0, state[0])
            pose_stamped.pose.orientation.x = quaternion[0]
            pose_stamped.pose.orientation.y = quaternion[1]
            pose_stamped.pose.orientation.z = quaternion[2]
            pose_stamped.pose.orientation.w = quaternion[3]

            path.poses.append(pose_stamped)
        return path

    def pathCallback(self, path: Path):
        # self.get_logger().info("path received")
        self.positions = [
            (pose.pose.position.x, pose.pose.position.y) for pose in path.poses
        ]
        if self.vec is None or len(path.poses) < 2 or len(self.vec) != len(path.poses):
            self.get_logger().info("update path, but path or vec is not ready, maybe too short or not set")
            return
        self.path.update(Spline(self.positions, self.vec, 0.02))

    def velocityCallback(self, msg: Float64MultiArray):
        self.vec = msg.data
        if self.positions == None or len(self.positions) < 2 or len(self.vec) != len(self.positions):
            self.get_logger().info("update vec, but path or vec is not ready, maybe too short or not set")
            return
        self.path.update(Spline(self.positions, self.vec, 0.02))

    # def pathCallback(self, path: Path, vec: Float64MultiArray):
    #    self.get_logger().info("path received")
    #    if len(path.poses) < 2 or len(vec.data) != len(path.poses):
    #        return
    #    positions = [
    #        (pose.pose.position.x, pose.pose.position.y) for pose in path.poses
    #    ]
    #    self.path.update(Spline(positions, vec.data, 0.02))


def main(args=None):
    rclpy.init(args=args)
    mpc_path_tracker = MpcPathTracker()
    rclpy.spin(mpc_path_tracker)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
